import paddle
from paddle.io import Dataset
import numpy as np


# 编写dataset类--与torch一样写就好了
class PoemData(Dataset):
    """
    构造诗歌数据集，继承paddle.io.Dataset
    参数:
        poems (list): 诗歌数据列表，每一个元素为一首诗歌，诗歌未经编码
        tokenizer: bert的分词模型
        max_len: 接收诗歌的最大长度
    """

    def __init__(self, poems, tokenizer, max_len=128):
        super(PoemData, self).__init__()
        self.poems = poems
        self.tokenizer = tokenizer
        self.max_len = max_len

    def __getitem__(self, idx):
        # 获取数据集的第idx行诗词的文本
        line = self.poems[idx]
        # print("line:",line)             # line: 还源去，欲求般若易。但息是非心，自然成大智。
        # 得到改行诗词文本的sent_id也就是数字编码
        token_line = self.tokenizer.encode(line)
        # print("token_line:", token_line)
        '''
        token_line: {
                    'input_ids': 
                    [101, 6820, 3975, 1343, 8024, 3617, 3724, 5663, 5735, 3211, 511, 852, 2622, 3221,
                     7478, 2552, 8024, 5632, 4197, 2768, 1920, 3255, 511, 102],
                     
                     'token_type_ids': [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
                     }
        '''
        '''
        Token的理解比较容易，是一串自然语言符号的最小语言单位（单词、词组、字符），将输入的一句话根据规定的最小语言单位划分开
        token_type当模型的目的是对句子进行分类或回答问题，此时一次输入模型的句子就不止一句，
            这里用它对一篇诗词进行短句处理
        Token Type IDs就是标识不同句子序列。通常"[CLS]"表示一个句子的开始，“[SEP]”表示一个句子的结束。
        '''
        token, token_type = token_line['input_ids'], token_line['token_type_ids']
        # print("token:",token, "token_type:",token_type)
        '''
        token: [101, 6820, 3975, 1343, 8024, 3617, 3724, 5663, 5735, 3211, 
        511, 852, 2622, 3221, 7478, 2552, 8024, 5632, 4197, 2768,
         1920, 3255, 511, 102] 
         token_type: [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
        '''
        if len(token) > self.max_len + 1:
            # print("token[-1:]:",token[-1:])
            # token[-1:]: [102] -- 结束符号[SEP]
            token = token[:self.max_len] + token[-1:]
            # print("token:",token)
            '''
            token: [101, 6820, 3975, 1343, 8024, 3617, 3724, 5663, 5735, 
                3211, 511, 852, 2622, 3221, 7478, 2552, 8024, 5632, 4197, 
                2768, 1920, 3255, 511, 102] 
            token_type: [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
            '''
            token_type = token_type[:self.max_len] + token_type[-1:]
        input_token, input_token_type = token[:-1], token_type[:-1]
        label_token = np.array((token[1:] + [0] * self.max_len)[:self.max_len], dtype='int64')
        # 输入填充，不够最大长度的都会进行填充0
        input_token = np.array((input_token + [0] * self.max_len)[:self.max_len], dtype='int64')
        input_token_type = np.array((input_token_type + [0] * self.max_len)[:self.max_len], dtype='int64')
        input_pad_mask = (input_token != 0).astype('float32')
        return input_token, input_token_type, input_pad_mask, label_token, input_pad_mask

    def __len__(self):
        return len(self.poems)

if __name__ == '__main__':
    from paddlenlp.transformers import BertTokenizer
    bert_tokenizer = BertTokenizer.from_pretrained('bert-base-chinese')
    PoemData = PoemData(
        poems=["松下中峰路，疎阴偃盖清。远云收海雨，数鴈过秋城。抚事心犹壮，卑栖气未平。近来唯此乐，高韵寄闲情。","还源去，欲求般若易。但息是非心，自然成大智。"],
        tokenizer=bert_tokenizer,
        max_len=25)
    print(PoemData.__len__())
    input_token1, input_token_type1, input_pad_mask1, label_token1, _=PoemData.__getitem__(1)
    print(input_token1)
    print("input_token1decode:",bert_tokenizer.decode(input_token1))
    print("label_token1decode:", bert_tokenizer.decode(label_token1))